import string

from collections import OrderedDict
import collections
## {{{ OrderedSet gratefully copied from http://code.activestate.com/recipes/576694/
KEY, PREV, NEXT = range(3)
class OrderedSet(collections.MutableSet):

	def __init__(self, iterable=None):
		self.end = end = [] 
		end += [None, end, end]		 # sentinel node for doubly linked list
		self.map = {}				   # key --> [key, prev, next]
		if iterable is not None:
			self |= iterable

	def __len__(self):
		return len(self.map)

	def __contains__(self, key):
		return key in self.map

	def add(self, key):
		if key not in self.map:
			end = self.end
			curr = end[PREV]
			curr[NEXT] = end[PREV] = self.map[key] = [key, curr, end]

	def discard(self, key):
		if key in self.map:		
			key, prev, next = self.map.pop(key)
			prev[NEXT] = next
			next[PREV] = prev

	def __iter__(self):
		end = self.end
		curr = end[NEXT]
		while curr is not end:
			yield curr[KEY]
			curr = curr[NEXT]

	def __reversed__(self):
		end = self.end
		curr = end[PREV]
		while curr is not end:
			yield curr[KEY]
			curr = curr[PREV]

	def pop(self, last=True):
		if not self:
			raise KeyError('set is empty')
		key = next(reversed(self)) if last else next(iter(self))
		self.discard(key)
		return key

	def __repr__(self):
		if not self:
			return '%s()' % (self.__class__.__name__,)
		return '%s(%r)' % (self.__class__.__name__, list(self))

	def __eq__(self, other):
		if isinstance(other, OrderedSet):
			return len(self) == len(other) and list(self) == list(other)
		return set(self) == set(other)

	def __del__(self):
		self.clear()					# remove circular references
## }}}


class AtomsiteCache :
	'''Class to take care of efficient access to atom_site data (efficient is a TODO ofcourse!). Use getCache and not constructor.
		To be used by classes in pdbents, not directly by end-user
		Needs OPTIMIZATION
	'''
	@staticmethod
	def getCache(datacon) :
		'''factory method, prevents duplication of caches with same data connector'''
		if not hasattr(AtomsiteCache, "cache") : AtomsiteCache.cache = {}
		if not AtomsiteCache.cache.has_key(datacon) : AtomsiteCache.cache[datacon] = AtomsiteCache(datacon)
		return AtomsiteCache.cache[datacon]
	def __init__(s, datacon) :
		s.dc = datacon
	def getOrderedResidueOccurrence(s) :
		'''return a list of residue occurrences in order in which they occur in atom_site.
			each item consists of (model,chain,resnum,icode,resname,altcode)
			this is needed sometimes for interpreting EDS results because model id is not reportted therein.'''
		orderedRes = OrderedSet()
		for adata in s.dc.getCategory('atom_site',None,'iter') :
			mid, cid = adata['pdbx_PDB_model_num'], adata['auth_asym_id']
			resname, resnum = adata['label_comp_id'], adata['auth_seq_id']
			try : icode = adata['pdbx_PDB_ins_code'] # optional in xml?!
			except KeyError : icode = ' '
			atomname, altcode = adata['label_atom_id'], adata['label_alt_id']
			orderedRes.add( (mid,cid,resnum,icode,resname,altcode) )
		return list(orderedRes)
	def getHierarchy(s) :
		'''return nested dictionaries model : chain : (resnum,icode) : resname : (atomname,altcode)'''
		s.prepareMCRA()
		return s.mcra
	def prepareMCRA(s) :
		'''make hieararchical structure for interesting bits of atomsite data - model : chain : (resnum,icode) : resname : (atomname,altcode)'''
		if not hasattr(s,'mcra') or s.mcra == None :
			s.mcra = OrderedDict() # TODO try weakref here?
			for adata in s.dc.getCategory('atom_site',None,'iter') :
				mid, cid = adata['pdbx_PDB_model_num'], adata['auth_asym_id']
				resname, resnum = adata['label_comp_id'], adata['auth_seq_id']
				try : icode = adata['pdbx_PDB_ins_code'] # optional in xml?!
				except KeyError : icode = ' '
				atomname, altcode = adata['label_atom_id'], adata['label_alt_id']
				if not s.mcra.has_key(mid) : s.mcra[mid] = OrderedDict()
				if not s.mcra[mid].has_key(cid) : s.mcra[mid][cid] = OrderedDict()
				if not s.mcra[mid][cid].has_key((resnum,icode)) : s.mcra[mid][cid][(resnum,icode)] = OrderedDict()
				if not s.mcra[mid][cid][(resnum,icode)].has_key(resname) : s.mcra[mid][cid][(resnum,icode)][resname] = OrderedSet()
				s.mcra[mid][cid][(resnum, icode)][resname].add((atomname, altcode))
	def getSubgroupCoords(s, mid, chid, resnum, icode, resname, altcode) :
		'''this is slow as it accesses atom_site records serially.... dont know how to make this faster at present.
		if altcode is different from '?','.',' ' then only it is checked.'''
		blankid = set([" ","?","."])
		retmap = OrderedDict()
		for adata in s.dc.getCategory('atom_site',None,'iter') :
			if mid != adata['pdbx_PDB_model_num'] : continue
			if chid != adata['auth_asym_id'] : continue
			if resnum != adata['auth_seq_id'] : continue
			if resname != adata['label_comp_id'] : continue
			if altcode != adata['label_alt_id'] and not set([altcode,adata['label_alt_id']]).issubset(blankid) : continue
			try : dc_icode = adata['pdbx_PDB_ins_code'] # optional in xml?!
			except KeyError : dc_icode = ' '
			if icode != dc_icode and not set([icode,dc_icode]).issubset(blankid) : continue
			retmap [adata['label_atom_id']] = tuple([ string.atof(adata['Cartn_'+xyz]) for xyz in ['x','y','z'] ])
		return retmap
	def prepareESSmapCRI(s) :
		'''map eid,said,seqid to a list of corresponding chid, resnum, icode.
		Note that thsi is done using atom_site table, so only modelled residues will figure.
		Maps to a list because in case of waters, seqid is same for all water groups...'''
		s.ess2cri = OrderedDict() # TODO try weakref here?
		for adata in s.dc.getCategory('atom_site',None,'iter') :
			eid, said, seqid = adata["label_entity_id"], adata["label_asym_id"], adata["label_seq_id"]
			cri = (adata["auth_asym_id"] , adata["auth_seq_id"] , adata["pdbx_PDB_ins_code"])
			if not s.ess2cri.has_key(eid) : s.ess2cri[eid] = OrderedDict()
			if not s.ess2cri[eid].has_key(said) : s.ess2cri[eid][said] = OrderedDict()
			if not s.ess2cri[eid][said].has_key(seqid) : s.ess2cri[eid][said][seqid] = OrderedSet()
			s.ess2cri[eid][said][seqid].add(cri)
	def getESSmapCRI(s, eid, said) :
		'''return list of chain,resnum,inscode given entity id, struct_asym id'''
		if not hasattr(s,'ess2cri') or s.ess2cri == None : s.prepareESSmapCRI()
		return s.ess2cri[eid][said].values()
	def getElementCounts(s) :
		'''return counts of atoms according to element types, for each entity, model, chain combination'''
		s.prepareElemCounts()
		return s.elemcounts
	def getZeroOccuCounts(s) :
		'''return counts of zero occu atoms for each entity, model, chain combo'''
		s.prepareElemCounts()
		return s.zerooccucounts
	def prepareElemCounts(s) :
		'''prepare map from eid,mid,chid to element counts. also count number of zero occu atoms.'''
		if not hasattr(s,'elemcounts') :
			counts = {} ; zerooccucounts = {}
			for adata in s.dc.getCategory('atom_site',None,'iter') :
				eid, mid, chid, elem = adata["label_entity_id"], adata["pdbx_PDB_model_num"], adata["auth_asym_id"], adata["type_symbol"]
				emc = (eid,mid,chid)
				if not counts.has_key(emc) : counts[emc] = {} ; zerooccucounts[emc] = 0
				if not counts[emc].has_key(elem) : counts[emc][elem] = 0
				counts[emc][elem] += 1
				if string.atof(adata["occupancy"]) < 1e-5 : zerooccucounts[emc] += 1
			s.elemcounts = counts
			s.zerooccucounts = zerooccucounts
	def prepareGroupCounts(s) :
		'''prepare counts for number of modelled groups per entity,mode,chain and how many have  alt conf'''
		if not hasattr(s,'groupcounts') :
			groupcounts = {}
			for adata in s.dc.getCategory('atom_site',None,'iter') :
				eid, mid, chid, resnum, icode, altcode = adata["label_entity_id"], adata["pdbx_PDB_model_num"], adata["auth_asym_id"], adata["auth_seq_id"], adata["pdbx_PDB_ins_code"], adata["label_alt_id"]
				emc = (eid, mid, chid)
				if not groupcounts.has_key(emc) : groupcounts[emc] = {'groups':set(),'altconf':set()}
				groupcounts[emc]['groups'].add( (eid,mid,chid,resnum,icode) )
				if not altcode in ['.','?',' '] : groupcounts[emc]['groups'].add( (eid,mid,chid,resnum,icode,altcode) )
			for emc in groupcounts.keys() :
				groupcounts[emc]['groups'] = len(groupcounts[emc]['groups'])
				groupcounts[emc]['altconf'] = len(groupcounts[emc]['altconf'])
			s.groupcounts = groupcounts
	def getNumModelledGroups(s, eid, mid, chid) :
		s.prepareGroupCounts()
		return s.groupcounts[(eid,mid,chid)]['groups']
	def getNumModelledGroupsWithAltconf(s, eid, mid, chid) :
		s.prepareGroupCounts()
		return s.groupcounts[(eid,mid,chid)]['altconf']
